Layernormgrad

计算 Layer Normalization 操作的梯度。该算子是 Layer Normalization 的反向传播部分，用于计算损失函数相对于输入 x、以及可学习参数 gamma 和 beta 的梯度。

\[ \begin{align}\begin{aligned}\text{dg}_i = \sum_{j} \text{dy}_j \cdot \frac{x_j - \mu}{\sqrt{\sigma^2 + \epsilon}}\\\text{db}_i = \sum_{j} \text{dy}_j\\\text{dx}_i = f(\text{dy}, x, \gamma, \mu, \sigma^2)\end{aligned}\end{align} \]

其中 :math:mu 是均值，:math:sigma^2 是方差，:math:epsilon 是一个为了防止除零而添加的极小值。dx 的计算较为复杂，它依赖于 dy、x 和 gamma。

输入：

  • x - 前向传播时的输入数据地址。

  • dy - 后续层反向传播回来的梯度数据地址。

  • var - 前向传播时计算出的方差（variance）地址。

  • mean - 前向传播时计算出的均值（mean）地址。

  • gamma - 前向传播时使用的可学习缩放参数 :math:gamma 地址。

  • param_num - 特征维度的大小，也是 gamma 和 beta 的大小。

  • param_size - 进行独立归一化的单元数量（例如批处理大小 Batch Size）。

  • block_num - 块的数量（通常等于 param_size）。

  • block_size - 每个块的大小（通常等于 param_num）。

  • core_mask - 核掩码（仅共享存储版本需要）。

输出：

  • dx - 计算出的关于输入 x 的梯度地址。

  • dg - 计算出的关于参数 gamma 的梯度地址。

  • db - 计算出的关于参数 beta 的梯度地址。

支持平台：

FT78NE MT7004

备注

  • FT78NE 支持fp32

  • MT7004 支持fp16, fp32

共享存储版本:

void hp_layernormgrad_s(half *x, half *dy, half *var, half *mean, half *gamma, int param_num, int param_size, int block_num, int block_size, half *dx, half *dg, half *db, int core_mask)
void fp_layernormgrad_s(float *x, float *dy, float *var, float *mean, float *gamma, int param_num, int param_size, int block_num, int block_size, float *dx, float *dg, float *db, int core_mask)

C调用示例：

 1//FT78NE示例
 2#include <stdio.h>
 3#include <layernormgrad.h> // 假设头文件名为 layernormgrad.h
 4
 5int main(int argc, char* argv[]) {
 6    // 假设在DDR空间
 7    float *x = (float *)0xA0000000;
 8    float *dy = (float *)0xA1000000;
 9    float *var = (float *)0xA2000000;
10    float *mean = (float *)0xA3000000;
11    float *gamma = (float *)0xA4000000;
12    float *dx = (float *)0xB0000000;
13    float *dg = (float *)0xB1000000;
14    float *db = (float *)0xB2000000;
15
16    int param_num = 256; // 特征维度
17    int param_size = 64; // Batch Size
18    int core_mask = 0xff;
19
20    fp_layernormgrad_s(x, dy, var, mean, gamma, param_num, param_size, param_size, param_num, dx, dg, db, core_mask);
21    return 0;
22}

私有存储版本:

void hp_layernormgrad_p(half *x, half *dy, half *var, half *mean, half *gamma, int param_num, int param_size, int block_num, int block_size, half *dx, half *dg, half *db)
void fp_layernormgrad_p(float *x, float *dy, float *var, float *mean, float *gamma, int param_num, int param_size, int block_num, int block_size, float *dx, float *dg, float *db)

C调用示例：

 1//FT78NE示例
 2#include <stdio.h>
 3#include <layernormgrad.h> // 假设头文件名为 layernormgrad.h
 4
 5int main(int argc, char* argv[]) {
 6    // 假设在L2空间
 7    float *x = (float *)0x10000000;
 8    float *dy = (float *)0x11000000;
 9    float *var = (float *)0x12000000;
10    float *mean = (float *)0x13000000;
11    float *gamma = (float *)0x14000000;
12    float *dx = (float *)0x15000000;
13    float *dg = (float *)0x16000000;
14    float *db = (float *)0x17000000;
15
16    int param_num = 256; // 特征维度
17    int param_size = 64; // Batch Size
18
19    fp_layernormgrad_p(x, dy, var, mean, gamma, param_num, param_size, param_size, param_num, dx, dg, db);
20    return 0;
21}